import os
import sherpa_onnx
import alsaaudio

PLAY_DEVICE = 'plughw:1,0,0'
alsa = alsaaudio.PCM(alsaaudio.PCM_PLAYBACK, alsaaudio.PCM_NORMAL, PLAY_DEVICE)
# Open 44100 Hz, mono, 16-bit little-endian PCM output
alsa.setchannels(1)
alsa.setrate(44100)
alsa.setformat(alsaaudio.PCM_FORMAT_S16_LE)

MELO_MODEL_PATH = "./models/tts/vits-melo-tts-zh_en"
SID = 0

LJS_MODEL_PATH = "./models/tts/vits-ljs"

FANCHEN_MODEL_PATH = "./models/tts/vits-zh-hf-fanchen-C"

cfgLJSTTS = sherpa_onnx.OfflineTtsConfig(
        model=sherpa_onnx.OfflineTtsModelConfig(
            vits=sherpa_onnx.OfflineTtsVitsModelConfig(
                model=LJS_MODEL_PATH + "/vits-ljs.onnx",
                lexicon=LJS_MODEL_PATH + "/lexicon.txt",
                tokens=LJS_MODEL_PATH + "/tokens.txt",
            ),
            provider="cpu",
            num_threads=4,
        ),
        max_num_sentences=3
    )    

cfgMeloTTS = sherpa_onnx.OfflineTtsConfig(
        model=sherpa_onnx.OfflineTtsModelConfig(
            vits=sherpa_onnx.OfflineTtsVitsModelConfig(
                model=MELO_MODEL_PATH + "/model.onnx",
                lexicon=MELO_MODEL_PATH + "/lexicon.txt",
                #data_dir=args.vits_data_dir,
                dict_dir=MELO_MODEL_PATH + "/dict",
                tokens=MELO_MODEL_PATH + "/tokens.txt",
            ),
            provider="cpu",
            num_threads=4,
        ),
        rule_fsts=MELO_MODEL_PATH + "/phone.fst," + MELO_MODEL_PATH + "/number.fst," + MELO_MODEL_PATH + "/new_heteronym.fst",
        max_num_sentences=3
    )

cfgFanchenTTS = sherpa_onnx.OfflineTtsConfig(
        model=sherpa_onnx.OfflineTtsModelConfig(
            vits=sherpa_onnx.OfflineTtsVitsModelConfig(
                model=FANCHEN_MODEL_PATH + "/vits-zh-hf-fanchen-C.onnx",
                lexicon=FANCHEN_MODEL_PATH + "/lexicon.txt",
                #data_dir=args.vits_data_dir,
                dict_dir=FANCHEN_MODEL_PATH + "/dict",
                tokens=FANCHEN_MODEL_PATH + "/tokens.txt",
            ),
            provider="cpu",
            num_threads=4,
        ),
        rule_fsts=FANCHEN_MODEL_PATH + "/phone.fst," + FANCHEN_MODEL_PATH + "/number.fst," + FANCHEN_MODEL_PATH + "/new_heteronym.fst",
        max_num_sentences=3
    )


tts = None


# amixer -c 1 set PCM 50%
def setVolume(volume):
    os.system(f"amixer -c 1 set PCM {volume}%")
    return f"Volume set to {volume}%"

setVolume(20)

def init(cfg = cfgMeloTTS):
    global tts
    tts = sherpa_onnx.OfflineTts(cfg)
    print("[tts] TTS loaded.")

def ttsSayText(text, speed=1.0, sid=0):
    audio = tts.generate(text, sid=sid, speed=speed)
    print(audio)
    samples = audio.samples # List of float
    # Encode to bytes (s16-le)
    samples_bytes = bytearray()
    for sample in samples:
        s16Sample = max(min(int(sample * 32767), 32767), -32768)
        samples_bytes += s16Sample.to_bytes(2, byteorder='little', signed=True)
    alsa.write(samples_bytes)


if __name__ == "__main__":
    ttsSayText("你好, 你的心情如何？")


